
import torch
from alike import ALike, configs
import cv2
# import glob
import os
from data.kitti import KITTI_IMG
from data.loaddata import LoadData
from torch.utils.data import DataLoader
import torch.nn as nn
from demo import SimpleTracker
import numpy as np
from utils.gen_color import create_pascal_label_colormap, to_color_img

# segmentation
datapath = {'kitti_path' : '/workspace/wzj/dataset/slam/kitti/data_odometry_color/sequences/00',
            'cityscape_path' : '/workspace/wzj/dataset/cityscape',

            # LCD
            'aachen_path' : '',
            'newcollege_path' : '/workspace/wzj/dataset/image_retrieval/New_College/Images',
            'oxford_path' : '/workspace/wzj/revisitop/data/datasets/roxford5k',
            'paris_path' : '/workspace/wzj/revisitop/data/datasets/rparis6k',
            'mapillary_path': '/workspace/wzj/dataset/image_retrieval/Mapillary'
}

eval_configs = {
    'alike-t': {'c1': 8, 'c2': 16, 'c3': 32, 'c4': 64, 'dim': 64, 'single_head': True, 'radius': 2,
                'model_path': os.path.join(os.path.split(__file__)[0], 'models', 'mapillary', 'KptSegGlb-non-const-lr.pth')},
    'alike-s': {'c1': 8, 'c2': 16, 'c3': 48, 'c4': 96, 'dim': 96, 'single_head': True, 'radius': 2,
                'model_path': os.path.join(os.path.split(__file__)[0], 'models', 'mapillary', 'KptSegGlb-non-const-lr.pth')},
    'alike-n': {'c1': 16, 'c2': 32, 'c3': 64, 'c4': 128, 'dim': 128, 'single_head': True, 'radius': 2,
                'model_path': os.path.join(os.path.split(__file__)[0], 'models', 'mapillary', 'KptSegGlb-non-const-lr.pth')},
    'alike-l': {'c1': 32, 'c2': 64, 'c3': 128, 'c4': 128, 'dim': 128, 'single_head': False, 'radius': 2,
                'model_path': os.path.join(os.path.split(__file__)[0], 'models', 'mapillary', 'KptSegGlb-non-const-lr.pth')},
}

def inference():
    # 推理设备
    device = torch.device("cuda: 1" if torch.cuda.is_available() else "cpu")
    # wandb.init(project="train_seg")
    # 定义模型
    model = ALike(**configs['alike-t'], 
                  device=device,
                  top_k=-1,
                  scores_th=0.2,
                  n_limit=5000)

    model = model.to("cuda")
    model.eval()